﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LPLib
{
    /// <summary>
    /// Provides tools for the translation of polynucleotide sequences into polypeptide sequences and the inverse.
    /// </summary>
    public static class Translator
    {
        /// <summary>
        /// Dictionary mapping codons to amino acids.
        /// </summary>
        static public Dictionary<string, char> CodonTable;

        /// <summary>
        /// Dictionary mapping amino acids to codons. Ambiguities are handled by using the IUPAC codes where possible, and by using the specialized codes
        /// 1 = first nucleotide in Leuceine codon
        /// 2 = last nucleotide in Leuceine codon
        /// 3 = first nucleotide in Arginine codon
        /// 4 = last nucleotide in Argeinine codon
        /// 5,6,7 = (ordered) nucleotides in serine codon
        /// </summary>
        static public Dictionary<char, string> ReverseCodonTable;

        static private Dictionary<char, int> encoder;

        static Translator()
        {
            initializeTranslators();
        }

        static private void initializeEncoder()
        {
            encoder = new Dictionary<char, int>();
            for (int i = 0; i < SymbolSets.AminoAcids.Length; i++)
            {
                encoder.Add(SymbolSets.AminoAcids[i], i);
            }
        }

        static private void initializeTranslators()
        {
            CodonTable = new Dictionary<string, char>();

            CodonTable.Add("AAA", 'K');
            CodonTable.Add("AAC", 'N');
            CodonTable.Add("AAG", 'K');
            CodonTable.Add("AAT", 'N');
            CodonTable.Add("ACA", 'T');
            CodonTable.Add("ACC", 'T');
            CodonTable.Add("ACG", 'T');
            CodonTable.Add("ACT", 'T');
            CodonTable.Add("AGA", 'R');
            CodonTable.Add("AGC", 'S');
            CodonTable.Add("AGG", 'R');
            CodonTable.Add("AGT", 'S');
            CodonTable.Add("ATA", 'I');
            CodonTable.Add("ATC", 'I');
            CodonTable.Add("ATG", 'M');
            CodonTable.Add("ATT", 'I');
            CodonTable.Add("CAA", 'Q');
            CodonTable.Add("CAC", 'H');
            CodonTable.Add("CAG", 'Q');
            CodonTable.Add("CAT", 'H');
            CodonTable.Add("CCA", 'P');
            CodonTable.Add("CCC", 'P');
            CodonTable.Add("CCG", 'P');
            CodonTable.Add("CCT", 'P');
            CodonTable.Add("CGA", 'R');
            CodonTable.Add("CGC", 'R');
            CodonTable.Add("CGG", 'R');
            CodonTable.Add("CGT", 'R');
            CodonTable.Add("CTA", 'L');
            CodonTable.Add("CTC", 'L');
            CodonTable.Add("CTG", 'L');
            CodonTable.Add("CTT", 'L');
            CodonTable.Add("GAA", 'E');
            CodonTable.Add("GAC", 'D');
            CodonTable.Add("GAG", 'E');
            CodonTable.Add("GAT", 'D');
            CodonTable.Add("GCA", 'A');
            CodonTable.Add("GCC", 'A');
            CodonTable.Add("GCG", 'A');
            CodonTable.Add("GCT", 'A');
            CodonTable.Add("GGA", 'G');
            CodonTable.Add("GGC", 'G');
            CodonTable.Add("GGG", 'G');
            CodonTable.Add("GGT", 'G');
            CodonTable.Add("GTA", 'V');
            CodonTable.Add("GTC", 'V');
            CodonTable.Add("GTG", 'V');
            CodonTable.Add("GTT", 'V');
            CodonTable.Add("TAA", '*');
            CodonTable.Add("TAC", 'Y');
            CodonTable.Add("TAG", '*');
            CodonTable.Add("TAT", 'Y');
            CodonTable.Add("TCA", 'S');
            CodonTable.Add("TCC", 'S');
            CodonTable.Add("TCG", 'S');
            CodonTable.Add("TCT", 'S');
            CodonTable.Add("TGA", '*');
            CodonTable.Add("TGC", 'C');
            CodonTable.Add("TGG", 'W');
            CodonTable.Add("TGT", 'C');
            CodonTable.Add("TTA", 'L');
            CodonTable.Add("TTC", 'F');
            CodonTable.Add("TTG", 'L');
            CodonTable.Add("TTT", 'F');

            ReverseCodonTable = new Dictionary<char, string>();
            // note the special codes:
            // 1 = first nucleotide in Leuceine codon
            // 2 = last nucleotide in Leuceine codon
            // 3 = first nucleotide in Arginine codon
            // 4 = last nucleotide in Argeinine codon
            // 5,6,7 = (ordered) nucleotides in serine codon
            ReverseCodonTable.Add('F', "TTY");
            ReverseCodonTable.Add('L', "1T2");
            ReverseCodonTable.Add('I', "ATH");
            ReverseCodonTable.Add('M', "ATG");
            ReverseCodonTable.Add('V', "GTN");
            ReverseCodonTable.Add('S', "567");
            ReverseCodonTable.Add('P', "CCN");
            ReverseCodonTable.Add('T', "CAN");
            ReverseCodonTable.Add('A', "GCN");
            ReverseCodonTable.Add('Y', "TAY");
            ReverseCodonTable.Add('H', "CAY");
            ReverseCodonTable.Add('Q', "CAR");
            ReverseCodonTable.Add('N', "AAY");
            ReverseCodonTable.Add('K', "AAR");
            ReverseCodonTable.Add('D', "GAY");
            ReverseCodonTable.Add('E', "GAR");
            ReverseCodonTable.Add('C', "TGY");
            ReverseCodonTable.Add('W', "TGG");
            ReverseCodonTable.Add('R', "3G4");
            ReverseCodonTable.Add('G', "GGN");
        }

        /// <summary>
        /// Performs the translation of a three-letter codon into a single-letter amino acid.
        /// </summary>
        /// <param name="s">The codon to be translated.</param>
        /// <returns>If s represents a valid codon, the correct single-letter amino acid code is returned. If not, returns "X".</returns>
        static private char translate(string s)
        {
            if (CodonTable.ContainsKey(s))
            {
                return CodonTable[s];
            }
            else return 'X';
        }

        /// <summary>
        /// Translates a polynucleotide string into a polypeptide string.
        /// </summary>
        /// <param name="s">The polynucleotide string.</param>
        /// <returns>The amino acid sequence that corresponds to the translation of s.</returns>
        static public string Translate(string s)
        {
            StringBuilder translation = new StringBuilder();
            for (int i = 0; i < s.Length - 2; i += 3)
            {
                translation.Append(translate(s.Substring(i, 3)));
            }
            return translation.ToString();
        }

        static public double[] Translate(double[][] codonpmf)
        {
            // codonpmf is an array comprising three independent nucleotide pmfs
            double[] apmf = new double[SymbolSets.AminoAcids.Length];
            for (int i0 = 0; i0 < 4; i0++ )
            {
                for (int i1 = 0; i1 < 4; i1++)
                {
                    for (int i2 = 0; i2 < 4; i2++)
                    {
                        string codon = SymbolSets.PureNucleotides[i0].ToString() + SymbolSets.PureNucleotides[i1].ToString() + SymbolSets.PureNucleotides[i2].ToString();
                        char aminoAcid = translate(codon);
                        int iaminoacid = encoder[aminoAcid];
                        apmf[iaminoacid] += codonpmf[0][i0] * codonpmf[1][i1] * codonpmf[2][i2];
                    }
                }
            }
            return apmf;
        }
    }
}


